# -*- coding: utf-8 -*-
"""
股票筛选器
根据以下规则筛选股票：
1. 当前股票必须多头排列，即10日线>20日线>60日线
2. 当前最新股价已经回踩20日线
3. 最近30个交易日的最高价不超过平均股价的10%
"""
import logging
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from database import DatabaseManager
from config import DATABASE_CONFIG
import warnings

# 忽略pandas的SettingWithCopyWarning警告
warnings.filterwarnings('ignore', category=pd.errors.SettingWithCopyWarning)

# 配置日志
# 获取项目根目录
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
LOG_DIR = os.path.join(PROJECT_ROOT, 'logs')
os.makedirs(LOG_DIR, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(os.path.join(LOG_DIR, 'stock_screener.log'), encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)

class StockScreener:
    """股票筛选器"""
    
    def __init__(self):
        self.db = DatabaseManager(DATABASE_CONFIG)
    
    def get_stock_daily_data(self, stock_code: str, days: int = 60) -> pd.DataFrame:
        """获取股票日线数据"""
        sql = """
        SELECT trade_date, open_price, high_price, low_price, close_price, volume, change_pct
        FROM new_stock_daily 
        WHERE stock_code = %s 
        ORDER BY trade_date DESC 
        LIMIT %s
        """
        try:
            results = self.db.execute_query(sql, (stock_code, days))
            if not results:
                return pd.DataFrame()
            
            df = pd.DataFrame(results)
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            
            # 转换价格和数值列为float类型，避免decimal类型计算问题
            price_columns = ['open_price', 'high_price', 'low_price', 'close_price']
            for col in price_columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')
            
            df['volume'] = pd.to_numeric(df['volume'], errors='coerce')
            df['change_pct'] = pd.to_numeric(df['change_pct'], errors='coerce')
            
            df = df.sort_values('trade_date')  # 按日期升序排列
            return df
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 数据失败: {e}")
            return pd.DataFrame()
    
    def calculate_ma(self, df: pd.DataFrame, periods: list = [10, 20, 60]) -> pd.DataFrame:
        """计算移动平均线"""
        if df.empty or len(df) < max(periods):
            return df
        
        # 确保close_price是float类型
        df['close_price'] = pd.to_numeric(df['close_price'], errors='coerce')
        
        for period in periods:
            df[f'ma{period}'] = df['close_price'].rolling(window=period).mean()
        
        return df
    
    def is_bullish_arrangement(self, df: pd.DataFrame) -> bool:
        """判断是否为多头排列（10日线 > 20日线 > 60日线）"""
        if df.empty or len(df) < 60:
            return False
        
        # 获取最新的移动平均线值
        latest = df.iloc[-1]
        
        # 检查是否有空值
        ma_columns = ['ma10', 'ma20', 'ma60']
        for col in ma_columns:
            if pd.isna(latest[col]):
                return False
        
        # 多头排列条件：10日线 > 20日线 > 60日线
        return latest['ma10'] > latest['ma20'] > latest['ma60']
    
    def is_price_near_ma20(self, df: pd.DataFrame, tolerance: float = 0.03) -> bool:
        """判断当前股价是否回踩20日线（允许3%的偏差）"""
        if df.empty or len(df) < 20:
            return False
        
        latest = df.iloc[-1]
        current_price = float(latest['close_price'])
        ma20 = float(latest['ma20'])
        
        if pd.isna(ma20) or ma20 == 0:
            return False
        
        # 计算股价与20日线的偏差百分比
        deviation = abs(current_price - ma20) / ma20
        
        # 如果偏差在容忍范围内，认为股价回踩20日线
        return deviation <= tolerance
    
    def is_price_stable(self, df: pd.DataFrame, days: int = 30, max_deviation: float = 0.10) -> bool:
        """判断最近N个交易日的最高价是否不超过平均股价的指定百分比"""
        if df.empty or len(df) < days:
            return False
        
        # 获取最近N个交易日的数据
        recent_data = df.tail(days)
        
        # 确保数据类型正确
        recent_data['close_price'] = pd.to_numeric(recent_data['close_price'], errors='coerce')
        recent_data['high_price'] = pd.to_numeric(recent_data['high_price'], errors='coerce')
        
        # 计算平均股价
        avg_price = recent_data['close_price'].mean()
        
        # 获取最高价
        max_price = recent_data['high_price'].max()
        
        if pd.isna(avg_price) or pd.isna(max_price) or avg_price == 0:
            return False
        
        # 计算最高价相对于平均价的偏差
        deviation = (max_price - avg_price) / avg_price
        
        # 如果偏差不超过指定百分比，认为价格稳定
        return deviation <= max_deviation
    
    def get_stock_info(self, stock_code: str) -> dict:
        """获取股票基本信息"""
        sql = "SELECT stock_name, exchange, industry, market FROM new_stock_info WHERE stock_code = %s"
        try:
            results = self.db.execute_query(sql, (stock_code,))
            return results[0] if results else {}
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 基本信息失败: {e}")
            return {}
    
    def get_all_stock_codes(self) -> list:
        """获取所有活跃股票代码"""
        sql = "SELECT stock_code FROM new_stock_info WHERE status = 'active'"
        try:
            results = self.db.execute_query(sql)
            return [row['stock_code'] for row in results]
        except Exception as e:
            logger.error(f"获取股票代码列表失败: {e}")
            return []
    
    def screen_stocks(self) -> pd.DataFrame:
        """筛选符合条件的股票"""
        logger.info("开始筛选股票...")
        
        stock_codes = self.get_all_stock_codes()
        if not stock_codes:
            logger.error("未找到任何股票代码")
            return pd.DataFrame()
        
        logger.info(f"共找到 {len(stock_codes)} 只股票，开始筛选...")
        
        screened_stocks = []
        processed = 0
        
        for stock_code in stock_codes:
            try:
                # 获取股票数据（需要60天数据来计算60日线）
                df = self.get_stock_daily_data(stock_code, 60)
                if df.empty:
                    continue
                
                # 计算移动平均线
                df = self.calculate_ma(df)
                
                # 检查条件1：多头排列
                if not self.is_bullish_arrangement(df):
                    processed += 1
                    if processed % 500 == 0:
                        logger.info(f"已处理 {processed}/{len(stock_codes)} 只股票")
                    continue
                
                # 检查条件2：股价回踩20日线
                if not self.is_price_near_ma20(df):
                    processed += 1
                    if processed % 500 == 0:
                        logger.info(f"已处理 {processed}/{len(stock_codes)} 只股票")
                    continue
                
                # 检查条件3：价格稳定
                if not self.is_price_stable(df):
                    processed += 1
                    if processed % 500 == 0:
                        logger.info(f"已处理 {processed}/{len(stock_codes)} 只股票")
                    continue
                
                # 所有条件都满足，添加到结果列表
                latest = df.iloc[-1]
                stock_info = self.get_stock_info(stock_code)
                
                # 计算额外指标
                recent_30 = df.tail(30)
                # 确保数据类型正确
                recent_30['close_price'] = pd.to_numeric(recent_30['close_price'], errors='coerce')
                recent_30['high_price'] = pd.to_numeric(recent_30['high_price'], errors='coerce')
                
                avg_price_30 = recent_30['close_price'].mean()
                max_price_30 = recent_30['high_price'].max()
                price_deviation = (max_price_30 - avg_price_30) / avg_price_30 * 100 if avg_price_30 > 0 else 0
                
                # 计算股价与20日线的偏差
                current_price = float(latest['close_price'])
                ma20 = float(latest['ma20'])
                price_ma20_deviation = (current_price - ma20) / ma20 * 100 if ma20 > 0 else 0
                
                screened_stocks.append({
                    'stock_code': stock_code,
                    'stock_name': stock_info.get('stock_name', ''),
                    'exchange': stock_info.get('exchange', ''),
                    'industry': stock_info.get('industry', ''),
                    'market': stock_info.get('market', ''),
                    'current_price': latest['close_price'],
                    'ma10': round(latest['ma10'], 2),
                    'ma20': round(latest['ma20'], 2),
                    'ma60': round(latest['ma60'], 2),
                    'change_pct': latest['change_pct'],
                    'volume': latest['volume'],
                    'trade_date': latest['trade_date'].strftime('%Y-%m-%d'),
                    'price_ma20_deviation_pct': round(price_ma20_deviation, 2),
                    'max_price_30': round(max_price_30, 2),
                    'avg_price_30': round(avg_price_30, 2),
                    'price_stability_pct': round(price_deviation, 2)
                })
                
                processed += 1
                if processed % 100 == 0:
                    logger.info(f"已处理 {processed}/{len(stock_codes)} 只股票，找到 {len(screened_stocks)} 只符合条件的股票")
                    
            except Exception as e:
                logger.error(f"筛选股票 {stock_code} 失败: {e}")
                processed += 1
                continue
        
        logger.info(f"筛选完成，共找到 {len(screened_stocks)} 只符合条件的股票")
        return pd.DataFrame(screened_stocks)
    
    def save_results(self, df: pd.DataFrame, filename: str = None):
        """保存筛选结果"""
        if df.empty:
            logger.warning("没有符合条件的股票数据可保存")
            return
        
        if filename is None:
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            filename = f'screened_stocks_{timestamp}.csv'
        
        try:
            df.to_csv(filename, index=False, encoding='utf-8-sig')
            logger.info(f"筛选结果已保存到: {filename}")
            
            # 同时保存股票代码文件
            stock_codes_filename = filename.replace('screened_stocks_', 'stock_codes_').replace('.csv', '.txt')
            stock_codes = ','.join(df['stock_code'].astype(str))
            with open(stock_codes_filename, 'w', encoding='utf-8') as f:
                f.write(stock_codes)
            logger.info(f"股票代码列表已保存到: {stock_codes_filename}")
            
        except Exception as e:
            logger.error(f"保存结果失败: {e}")
    
    def print_results(self, df: pd.DataFrame):
        """打印筛选结果"""
        if df.empty:
            print("未找到符合条件的股票")
            return
        
        print(f"\n=== 股票筛选结果 ===")
        print(f"筛选条件：")
        print(f"  1. 多头排列：10日线 > 20日线 > 60日线")
        print(f"  2. 回踩20日线：当前股价接近20日线（偏差≤3%）")
        print(f"  3. 价格稳定：最近30日最高价不超过平均价10%")
        print(f"共找到 {len(df)} 只符合条件的股票")
        print(f"筛选时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print("\n股票列表:")
        print("-" * 150)
        print(f"{'代码':<8} {'名称':<12} {'交易所':<6} {'行业':<20} {'当前价':<8} {'MA10':<8} {'MA20':<8} {'MA60':<8} {'回踩%':<8} {'稳定性%':<8} {'涨跌幅':<8} {'日期':<12}")
        print("-" * 150)
        
        for _, row in df.iterrows():
            print(f"{row['stock_code']:<8} {row['stock_name']:<12} {row['exchange']:<6} "
                  f"{str(row['industry'])[:18]:<20} {row['current_price']:<8.2f} "
                  f"{row['ma10']:<8.2f} {row['ma20']:<8.2f} {row['ma60']:<8.2f} "
                  f"{row['price_ma20_deviation_pct']:<8.2f} {row['price_stability_pct']:<8.2f} "
                  f"{row['change_pct']:<8.2f} {row['trade_date']:<12}")
        
        print("-" * 150)
        
        # 按行业统计
        if 'industry' in df.columns:
            industry_stats = df['industry'].value_counts()
            print(f"\n按行业分布:")
            for industry, count in industry_stats.head(10).items():
                print(f"  {industry}: {count} 只")
        
        # 按交易所统计
        if 'exchange' in df.columns:
            exchange_stats = df['exchange'].value_counts()
            print(f"\n按交易所分布:")
            for exchange, count in exchange_stats.items():
                print(f"  {exchange}: {count} 只")
        
        # 统计信息
        print(f"\n统计信息:")
        print(f"  平均股价: {df['current_price'].mean():.2f}")
        print(f"  平均回踩偏差: {df['price_ma20_deviation_pct'].mean():.2f}%")
        print(f"  平均价格稳定性: {df['price_stability_pct'].mean():.2f}%")
        print(f"  平均涨跌幅: {df['change_pct'].mean():.2f}%")

def main():
    """主函数"""
    try:
        screener = StockScreener()
        
        # 筛选股票
        screened_df = screener.screen_stocks()
        
        # 打印结果
        screener.print_results(screened_df)
        
        # 保存结果
        screener.save_results(screened_df)
        
    except Exception as e:
        logger.error(f"筛选过程出错: {e}")
        raise

if __name__ == "__main__":
    main()
